"""
机械臂平衡任务配置

定义机械臂倒立摆平衡任务的配置和环境创建。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import gymnasium as gym

from parnassus_train.envs import GrpcArmEnv


@dataclass
class ArmBalanceTask:
    """
    机械臂倒立摆平衡任务配置
    
    该任务的目标是让机械臂的摆杆保持竖直向上的平衡状态。
    """
    
    # 环境配置
    server_address: str = "localhost:50051"
    obs_dim: int = 4
    action_dim: int = 1
    max_episode_steps: int = 500
    reward_shaping: bool = True
    
    # 训练配置
    num_episodes: int = 10000
    learning_rate: float = 3e-4
    gamma: float = 0.99
    gae_lambda: float = 0.95
    
    # PPO 配置
    clip_epsilon: float = 0.2
    value_coef: float = 0.5
    entropy_coef: float = 0.01
    max_grad_norm: float = 0.5
    n_epochs: int = 10
    batch_size: int = 64
    
    # 模型配置
    hidden_dim: int = 64
    
    # 日志和保存
    log_interval: int = 10
    save_interval: int = 100
    save_dir: str = "checkpoints"
    use_wandb: bool = True
    wandb_project: str = "parnassus-ppo"
    wandb_name: Optional[str] = None
    
    def create_env(self) -> gym.Env:
        """
        创建环境
        
        Returns:
            env: Gymnasium 环境实例
        """
        env = GrpcArmEnv(
            server_address=self.server_address,
            obs_dim=self.obs_dim,
            action_dim=self.action_dim,
            max_episode_steps=self.max_episode_steps,
            reward_shaping=self.reward_shaping,
        )
        return env
    
    def to_dict(self) -> dict:
        """
        转换为字典（用于配置记录）
        
        Returns:
            config_dict: 配置字典
        """
        return {
            "server_address": self.server_address,
            "obs_dim": self.obs_dim,
            "action_dim": self.action_dim,
            "max_episode_steps": self.max_episode_steps,
            "reward_shaping": self.reward_shaping,
            "num_episodes": self.num_episodes,
            "learning_rate": self.learning_rate,
            "gamma": self.gamma,
            "gae_lambda": self.gae_lambda,
            "clip_epsilon": self.clip_epsilon,
            "value_coef": self.value_coef,
            "entropy_coef": self.entropy_coef,
            "max_grad_norm": self.max_grad_norm,
            "n_epochs": self.n_epochs,
            "batch_size": self.batch_size,
            "hidden_dim": self.hidden_dim,
            "log_interval": self.log_interval,
            "save_interval": self.save_interval,
            "save_dir": self.save_dir,
            "use_wandb": self.use_wandb,
            "wandb_project": self.wandb_project,
            "wandb_name": self.wandb_name,
        }
